library(Matrix)

# Main model structures 
setClass("mm",
 representation(          
   call = "call",
   X = "dgCMatrix",
   y = "numeric",
   dims = "integer",           
   eta = "numeric",
   mu = "numeric",
   resid = "numeric",
   wts = "numeric",
   beta = "matrix",
   pi = "matrix",
   control = "numeric",
   subject = "integer", 
   gid = "integer",
   gvar = "integer",
   frame = "data.frame",
   formula = "formula")
)


# the nrmm class for the non-regularized standard multinomial logit model
setClass("nrmm", 
 representation(                        
   hess = "matrix",
   L = "CHMfactor"),
 contains = "mm"
)

# the rmm class for regularized multinomial logit model
setClass("rmm", 
 representation(                   
   lambda = "numeric",
   pfct = "numeric"),
 contains = "mm"
)

# the group lasso class
setClass("glmm", 
 contains = "rmm"
)

# validating the model struct
val.mm <- function(object){
  dims <- object@dims
  if ((nrow(object@X) != dims["nO"]) || (ncol(object@X) != dims["nB"]))
    stop("Incorrect dimensions of 'X'.")
  # y can be null in prediction  
  #  if ((nrow(object@y) != dims["n"]) || (ncol(object@y) != 1))
  #    stop("Incorrect dimensions of 'y'.")
  if (length(object@eta) != dims["nO"])
    stop("Incorrect length of 'eta'.")
  if (length(object@mu) != dims["nO"] || any(object@mu < 0))
    stop("Invalid 'mu'.")
  if (length(object@resid) != dims["nO"])
    stop("Incorrect length of 'resid'.")
  if (length(object@wts) != dims["nO"] || any(object@wts < 0))
    stop("Invalid 'wts'.")    
  if (nrow(object@beta) != dims["nB"] * dims["nS"] || 
        ncol(object@beta) != dims["nL"])
    stop("Incorrect dimension of 'beta'.")
  if (nrow(object@pi) != dims["nS"] || 
        ncol(object@pi) != dims["nL"])
    stop("Incorrect dimension of 'pi'.")
 # if (length(object@cset) != dims["nO"] || any(!(object@cset %in% c(0, 1))))
 #    stop("Invalid 'cset'.")
  return(TRUE)
}

# initialize the model struct
setMethod("initialize", "mm",
  function(.Object, ...) {
    .Object <- callNextMethod()
    dims <- .Object@dims
    nm <- c("nO", "nB", "nP", "nC", "nG", "nL", "nI", "nS")  
    mt <- nm %in% names(dims)
    if (sum(mt) < length(nm) || any(dims <= 0))
      stop("Invalid 'dims' slot.")      
#    if (dims["nO"] != dims["nP"] * dims["nC"])
#      stop("Incorrect dimensions in 'dims'.")      
    if (length(.Object@eta) == 0)
      .Object@eta <- double(dims["nO"])    
    if (length(.Object@mu) == 0)
      .Object@mu <- double(dims["nO"])
    if (length(.Object@resid) == 0)
      .Object@resid <- double(dims["nO"])    
    if (length(.Object@wts) == 0)
      .Object@wts <- as.double(rep(1, dims["nO"]))        
    if (length(.Object@call) == 0)
      .Object@call <- call("foo") 
    if (length(.Object@frame) == 0)
      .Object@frame <- data.frame(list())    
    if (length(.Object@control) == 0)
      .Object@control <- rmm.control()
    if (length(.Object@beta) == 0){
      nn <- dims["nB"] * dims["nS"] * dims["nL"]
      .Object@beta <- matrix(rnorm(nn, sd = 0.01), ncol = dims["nL"])
    }
    if (length(.Object@pi) == 0)
      .Object@pi <- matrix(1/dims["nS"], dims["nS"], ncol = dims["nL"])
  #  if (length(.Object@cset) == 0)
  #    .Object@cset <- as.integer(dims["nO"])   
    val.mm(.Object)
    return(.Object)
  })


# validating the nrmm object
val.nrmm <- function(object){
  dims <- object@dims
  ctrl <- object@control  
  if ((nrow(object@hess) != dims["nB"]) || (ncol(object@hess) != dims["nB"]))
    stop("Incorrect dimensions of 'hess'.")
  if (ctrl["newton"] == 0 && dim(object@L) == 0)
    stop("Slot 'L' cannot be empty when using the 'MM' update.")
  return(TRUE)
}

# initialize the nrmm object
setMethod("initialize", "nrmm",
  function(.Object, ...) {
    .Object <- callNextMethod()
    dims <- .Object@dims
    if (length(.Object@hess) == 0)
      .Object@hess <- matrix(0, dims["nB"], dims["nB"])             
    # no initialize method for class CHMfactor
    #if (length(.Object@L) == 0)
    #  .Ojbect@L <- new("dCHMsimpl")    
    val.nrmm(.Object)
    return(.Object)
  })


# validating the rmm object
val.rmm <- function(object){
  dims <- object@dims
  if (length(object@lambda) != dims["nL"] || any(object@lambda < 0))
    stop("Incompatible 'lambda'.")
  # assuming nG = nB for non-grouplasso methods
  if (length(object@pfct) != dims["nG"] || any(object@pfct < 0))
    stop("Incompatible 'pfct'.") 
  if (length(object@gvar) != dims["nG"] + 1 || any(object@gvar < 0))
    stop("Incompatible 'gvar'.")
  #if (length(object@ncols) != dims["nG"] || any(object@ncols <= 0))
  #  stop("Incompatible 'ncols'.")  
  return(TRUE)
}

# initialize the rmm object
setMethod("initialize", "rmm",
  function(.Object, ...) {
    .Object <- callNextMethod()
    dims <- .Object@dims
    if (length(.Object@pfct) == 0)
      .Object@pfct <- as.numeric(rep(1, dims["nG"]))    
    val.rmm(.Object)
    return(.Object)
  })



setMethod("coef", signature(object = "mm"),
  function(object, which.lambda = 1, scale = TRUE, ...){
    control <- object@control
    dims <- object@dims
    gvar <- object@gvar
    nG <- dims["nG"]
    beta <- object@beta[, which.lambda]
    pi <- as.vector(object@pi[, which.lambda])
    od <- order(pi, decreasing = TRUE)
    beta <- rbind(matrix(beta, ncol = dims["nS"]), pi)
    beta <- beta[, od, drop = FALSE]    
    # scale back to the original scale if necessary
    if (control["standardize"] && scale){
      U <- attr(object@X, "U")
      for (i in 1:nG){
        idx <- (gvar[i] + 1):gvar[i + 1]
        beta[idx, ] <- as.numeric(solve(U[[i]]) %*% beta[idx, , drop = FALSE])
      }
    }    
    nm <- c(colnames(object@X), "pi")
    dimnames(beta) <- list(nm, as.character(1:dims["nS"]))
    # remove pi if nrmm
    if (dims["nS"] == 1)
      beta <- beta[-nrow(beta), , drop = FALSE]
    return(beta)
  }
)

setMethod("fitted", signature(object = "mm"),
  function(object, which.lambda = 1, ...){
    nS <- object@dims["nS"]
    nB <- object@dims["nB"]
    dims2 <- object@dims
    dims2[c("nL", "nS")] <- as.integer(c(1, 1))
    ans <- new("mm", X = object@X, y = object@y, dims = dims2,
               formula = object@formula, gid = object@gid, 
               control = object@control)   
    mu <- matrix(0, dims2["nO"], nS)
    for (i in 1:nS){
      idx <- (nB * (i - 1) + 1):(nB * i)
      ans@beta <- matrix(object@beta[idx, which.lambda])
      .Call("R_rmm_update_prob", ans)
      mu[, i] <- ans@mu
    }
    row.names(mu) <- row.names(object@X)
    colnames(mu) <- as.character(1:nS)  
    return(mu)
  }
)

setMethod("formula", signature(x = "mm"),
  function (x, ...) 
    x@formula
)

#FIXME: do we need this?
setMethod("update", signature(object = "mm"),
  function (object, formula., evaluate = TRUE, ...) {
    if (is.null(call <- getCall(object))) 
      stop("need an object with call component")
    extras <- match.call(expand.dots = FALSE)$...
    if (!missing(formula.)) 
      call$formula <- update.formula(formula(object), formula.)
    if (length(extras)) {
      existing <- !is.na(match(names(extras), names(call)))
      for (a in names(extras)[existing]) call[[a]] <- extras[[a]]
      if (any(!existing)) {
        call <- c(as.list(call), extras[!existing])
        call <- as.call(call)
      }
    }
    if (evaluate) 
      eval(call, parent.frame())
    else call
  }
)


#'
#' Generate predicted values given an new input data set
#' and a fitted hlogit model. 
#'
#' @param object a fitted hlogit object
#' @param newdata a data frame containing the necessary predictors for predictions
#' @param type whether the predicted probability or outcome is to be returned      
#'
#' @return  a list of predicted values (length n) for each value of rho
#'    
#' @note The steps involved in this function are the following:
#' 1. generate design matrix for the new data
#' 3. call R_update_prob to compute the estimated probability 
#'    based on the estimated beta and predicted gamma    

setMethod("predict", signature(object = "mm"),
function(object, newdata, which.lambda = 1, save.mm = TRUE, ...) {
  #check.data(newdata) # check newdata
  if (which.lambda < 0 || which.lambda > object@dims["nL"])
    stop("'which.lambda' is out of range.")
  if (length(which.lambda) > 1){
    which.lambda <- which.lambda[1]
    warning("Only the first element of 'which.lambda' is used.")
  }
  # construct the new design matrix
  mc <- object@call
  mc$data <- newdata
  tt <- attr(object@frame, "terms")  
  Terms <- object@formula
  resp <- names(attr(tt, "dataClasses"))[attr(tt, "response")]
  if (!(resp %in% names(newdata)))
    Terms <- delete.response(tt)
  xlevels <- .getXlevels(tt, object@frame)
  fr <- getFrame(mc, Terms, object@control["use.contrast"], 
                 xlev = xlevels, drop.unused.levels = FALSE)
  X <- fr$X
  gvar <- object@gvar
  # apply scaling 
  if (object@control["standardize"]){
    nG <- length(gvar) - 1
    U <- attr(object@X, "U")
    X2 <- Matrix(0, nrow(X), 0)
    for (i in 1:nG){
      idx <- (gvar[i] + 1):gvar[i + 1]
      XX <- X[, idx, drop = FALSE]
      X2 <- cBind(X2, XX %*% solve(U[[i]]))
    }
    dimnames(X2) <- dimnames(X)
    X <- X2
  }
  # update dims and control
  dims <- object@dims  
  dims[c("nO", "nP", "nC", "nI", "nL")] <- 
    as.integer(c(nrow(X), nlevels(fr$chid), nlevels(fr$alt), max(fr$id) + 1, 1))
      
  # generate predictions
  beta <- object@beta[, which.lambda, drop = FALSE]
  pi <- object@pi[, which.lambda, drop = FALSE]
  
  tmp <- new("mm", X = X, y = fr$y, dims = dims, control = object@control, 
             formula = object@formula, pi = pi, subject = fr$id,
             wts = fr$wts, gid = fr$gid)
  mu <- vector("list", dims["nS"])
  for (i in 1:dims["nS"]){
    tmp@beta <- matrix(beta[(1 + dims["nB"] * (i - 1)):(dims["nB"] * i)])
    .Call("R_rmm_update_prob", tmp)
    mu[[i]] <- tmp@mu
  }
  mu <- do.call(cbind, mu) %*% pi
  #yp <- apply(matrix(mu, nrow = dims["nC"]), 2, which.max)
  # return predictions
  tmp@beta <- beta
  list(mm = if(save.mm) tmp else NULL, mean = mu)
})


# subset method for the "mlogit.data" class
# this takes care of the subsetting of the index 
"[.mlogit.data" <- function(x, i, j, drop = FALSE, ...) {
  xx <- NextMethod()
  ix <- index(x)  # this may be null in mlogit
  if (!missing(i) & !is.null(ix)){    
    ixx <- ix[i, , FALSE]
    ix <- droplevels(ixx)
  } 
  attr(xx, "index") <- ix
  xx
}

index.data.frame <- function (x, ...) 
{
  attr(x, "index")
}